import pylab;

class GraphSceneModel(object):
    def __init__(self, objGraph):
        self.objGraph = objGraph;
    
    []#Set Graph Values
    def setNodes_GrpActsInfo(self, nodes, dfGrpActs):
        for i in range(len(nodes)):
            'Activity', 'Start-Frm', 'End-Frm', 'Role', 'Persons'
            nodes[i].nodeModel.setGrpId(i);
            nodes[i].nodeModel.setAct(dfGrpActs["Activity"][i]);
            nodes[i].nodeModel.setStartFrm(dfGrpActs["Start-Frm"][i]);
            nodes[i].nodeModel.setEndFrm(dfGrpActs["End-Frm"][i]);
            nodes[i].nodeModel.setRole(dfGrpActs["Role"][i]);
            nodes[i].nodeModel.setPersons(str(dfGrpActs["Persons"][i]));                
    
    def setNodes_StateId(self, nodes, stateId):
        pass;
    
    def setNodes_StateDim(self, nodes, aStateDim):
        for i in range(len(nodes)):
            nodes[i].nodeModel.setStateDim(aStateDim[i]);
    
    def setNodes_Value(self, nodes, aValues):
        for i in range(len(nodes)):
            nodes[i].nodeModel.setValue(aValues[i]);
    
    def setNodes_Distance(self, nodes, aDists):
        for i in range(len(nodes)):
            nodes[i].nodeModel.setDistance(aDists[i]);       
        
    def setNodes_Reward(self, nodes, aRewards):
        for i in range(len(nodes)):
            nodes[i].nodeModel.setReward(aRewards[i]);
        
    def setNodes_StateType(self, nodes, aTxtStType):
        for i in range(len(nodes)):            
            if aTxtStType[i] == "Normal":
                nodes[i].nodeModel.setStateType(0);
            elif aTxtStType[i] == "Current":
                nodes[i].nodeModel.setStateType(1);
            elif aTxtStType[i] == "Fatal":
                nodes[i].nodeModel.setStateType(2);
            elif aTxtStType[i] == "Goal":
                nodes[i].nodeModel.setStateType(3);
    
    def setActions_QVals(self, nodes, mQVals):
        for i in range(len(nodes)):
            for j in range(len(nodes[i].actions)):
                nodes[i].actions[j].actModel.setQValue(mQVals[i][j]);
    
    def setActions_ActType(self, nodes, aTxtActType, aPolicy=None):
        for i in range(len(nodes)):
            for j in range(len(nodes[i].actions)):
                if aTxtActType[i][j] == "Normal" or aTxtActType[i][j] == "Good Action":
                    nodes[i].actions[j].actModel.setActionType(0);
                elif aTxtActType[i][j] == "Current":
                    nodes[i].actions[j].actModel.setActionType(1);
                elif aTxtActType[i][j] == "Risky" or aTxtActType[i][j] == "Risky Action":
                    nodes[i].actions[j].actModel.setActionType(2);
                elif aTxtActType[i][j] == "Best":
                    nodes[i].actions[j].actModel.setActionType(3);
            if aPolicy!=None and aPolicy[i]>0 and aPolicy[i]<=len(nodes[i].actions):
                k = aPolicy[i]-1;
                if aTxtActType[i][k]!="Risky" and aTxtActType[i][k]!="Risky Action":
                    nodes[i].actions[k].actModel.setActionType(3); #Best
    
    def setEdges_Trans(self, edges, trans):
        pass;
    
        
    
    []#Get State Values
    def getNodes_StateId(self, nodes):
        aStId = [];
        for node in nodes:
            aStId.append(node.nodeModel.getStateId());
        return aStId;
    
    def getNodes_StateDim(self, nodes):
        aStDim = [];
        for node in nodes:
            aStDim.append(node.nodeModel.getStateDim());
        return aStDim;
    
    def getNodes_Value(self, nodes):
        aVal = [];
        for node in nodes:
            aVal.append(node.nodeModel.getValue());
        return aVal;
        
    def getNodes_Reward(self, nodes):
        aReward = [];
        for node in nodes:
            aReward.append(node.nodeModel.getReward());
        return aReward;
    
    def getEdges_TransVals(self, edges):
        nStates = len(self.objGraph.nodes);
        if nStates > 0:
            nActs = len(self.objGraph.nodes[0].actions);
            aTrans = pylab.zeros((nActs, nStates,nStates));
        for edge in edges:
            nActFrom = edge.objFrom.actModel.getActId()-1;
            nStFrom = edge.objFrom.parentNode.nodeModel.getStateId()-1;
            nStTo = edge.objTo.nodeModel.getStateId()-1;
            nProb = edge.edgeModel.getTransitionValue();
            aTrans[nActFrom][nStFrom][nStTo] = nProb;
        return aTrans;
    
    def getEdges_Info(self, edges):
        aEdgeInfo = [];
        for edge in edges:
            nActFrom = edge.objFrom.actModel.getActId();
            nStFrom = edge.objFrom.parentNode.nodeModel.getStateId();
            nStTo = edge.objTo.nodeModel.getStateId();
            nProb = edge.edgeModel.getTransitionValue();
            aEdgeInfo.append([nActFrom, nStFrom, nStTo, nProb]);            
        return aEdgeInfo;   
    
    def getActions_QVals(self, nodes):
        aQVals = [];
        for node in nodes:
            aStQVals = [];
            for action in node.actions:
                aStQVals.append(action.actModel.getQValue());
            aQVals.append(aStQVals);
        return aQVals;
